#include "stdafx.h"

#include <qpdf/QPDF.hh>
#include <qpdf/QUtil.hh>
#include <qpdf/QTC.hh>
#include <qpdf/Pl_StdioFile.hh>
#include <qpdf/QPDFExc.hh>
#include <qpdf/QPDFWriter.hh>

#include "PdfProtector.h"

HINSTANCE g_hInstance = 0;
long      g_cLocks    = 0;
const WCHAR wzKey[] = L"I am protected";
const WCHAR wzSignedIL[] = L"SignedIL";
const WCHAR wzServerID[] = L"ServerID"; 
const WCHAR wzLicenses[] = L"Licenses";
const WCHAR wzRightsTemplate[] = L"RightsTemplate";
const WCHAR wzListGuid[] = L"ListGuid";
const WCHAR wzContent[] = L"Content";
const WCHAR wzRightsMask[] = L"RightsMask";
const WCHAR wzRequestingUser[] = L"RequestingUser";
const WCHAR wzURL[] = L"URL";
const WCHAR wzPolicyTitle[] = L"PolicyTitle";
const WCHAR wzPolicyDescription[] = L"PolicyDescription";
const WCHAR wzOfflineDays[] = L"OfflineDays";
const int STREAM_MAX = 32;
const WCHAR wzEULPrefix[] = L"EUL-";
#define STACK_RW_BUF_SIZE (8192)

#ifndef Unreferenced
#define Unreferenced(x)    ((void)x)
#endif

#ifndef cElements
#define cElements(ary) (sizeof(ary) / sizeof(ary[0]))
#endif

#define GUIDLEN 38

HRESULT HrEnsureStg(IStorage *pistg, const WCHAR *wzStg, bool fReadWrite, IStorage **ppistg)
{
  HRESULT hr = S_OK;

  if (pistg == NULL || wzStg == NULL || ppistg == NULL)
    return E_INVALIDARG;

  *ppistg = NULL;
  hr = pistg->OpenStorage(wzStg, NULL, (fReadWrite ? STGM_READWRITE : 0) | STGM_SHARE_EXCLUSIVE, NULL, 0, ppistg);
  if (hr == STG_E_FILENOTFOUND && fReadWrite)
    hr = pistg->CreateStorage(wzStg, STGM_CREATE|STGM_READWRITE|STGM_SHARE_EXCLUSIVE, 0, 0, ppistg);

  if (FAILED(hr))
  {
    if (*ppistg != NULL)
    {
      (*ppistg)->Release();
      *ppistg = NULL;
    }
  }
  return hr;
}

HRESULT HrEnsureStm(IStorage *pistg, const WCHAR *wzStm, bool fReadWrite, IStream **ppistm)
{
  HRESULT hr = S_OK;

  if (pistg == NULL || wzStm == NULL || ppistm == NULL)
    return E_INVALIDARG;

  *ppistm = NULL;
  hr = pistg->OpenStream(wzStm, NULL, (fReadWrite ? STGM_READWRITE : 0) | STGM_SHARE_EXCLUSIVE, 0, ppistm);
  if (hr == STG_E_FILENOTFOUND && fReadWrite)
    hr = pistg->CreateStream(wzStm, STGM_CREATE|STGM_READWRITE|STGM_SHARE_EXCLUSIVE, 0, 0, ppistm);

  if (FAILED(hr))
  {
    if (*ppistm != NULL)
    {
      (*ppistm)->Release();
      *ppistm = NULL;
    }
  }
  return hr;
}

HRESULT HrWriteRaw(IStream *pistm, const WCHAR *wzData)
{
  HRESULT hr = S_OK;
  DWORD cbSrc = 0;
  DWORD cbWritten = 0;

  if (pistm == NULL || wzData == NULL)
    return E_INVALIDARG;

  cbSrc = ((DWORD)wcslen(wzData) + 1) * sizeof(WCHAR);
  hr = pistm->Write((void*)&cbSrc, sizeof(DWORD), &cbWritten);
  if (SUCCEEDED(hr) && sizeof(DWORD) != cbWritten)
    hr = STG_E_WRITEFAULT;
  if (FAILED(hr))
    goto LExit;
  if (cbSrc > 0)
  {
    hr = pistm->Write((void*)wzData, cbSrc, &cbWritten);
    if (SUCCEEDED(hr) && cbSrc != cbWritten)
      hr = STG_E_WRITEFAULT;
    if (FAILED(hr))
      goto LExit;
  }

LExit:
  return hr;
}

HRESULT HrWriteRaw(IStream *pistm, const DWORD dwData)
{
  HRESULT hr = S_OK;
  DWORD cbSrc = 0;
  DWORD cbWritten = 0;

  if (pistm == NULL || dwData == NULL)
    return E_INVALIDARG;

  cbSrc = sizeof(DWORD);
  hr = pistm->Write((void*)&cbSrc, sizeof(DWORD), &cbWritten);
  if (SUCCEEDED(hr) && sizeof(DWORD) != cbWritten)
    hr = STG_E_WRITEFAULT;
  if (FAILED(hr))
    goto LExit;
  if (cbSrc > 0)
  {
    hr = pistm->Write((void*)&dwData, cbSrc, &cbWritten);
    if (SUCCEEDED(hr) && cbSrc != cbWritten)
      hr = STG_E_WRITEFAULT;
    if (FAILED(hr))
      goto LExit;
  }

LExit:
  return hr;
}

HRESULT HrReadRaw(IStream *pistm, WCHAR **pwz)
{
  HRESULT hr = S_OK;
  DWORD cbSrc = 0;
  DWORD cbRead = 0;

  if (pistm == NULL || pwz == NULL)
  {		
    return E_INVALIDARG;
  }

  *pwz = NULL;

  hr = pistm->Read((void*)&cbSrc, sizeof(DWORD), &cbRead);
  if (SUCCEEDED(hr) && sizeof(DWORD) != cbRead)
    hr = STG_E_READFAULT;
  if (FAILED(hr))
    goto LExit;

  *pwz = (WCHAR *)new BYTE [cbSrc];
  if (*pwz == NULL)
  {
    hr = E_OUTOFMEMORY;
    goto LExit;
  }

  hr = pistm->Read(*pwz, cbSrc, NULL);
  if (FAILED(hr))
    goto LExit;

LExit:
  if (FAILED(hr) && *pwz != NULL)
  {
    delete [] *pwz;
    *pwz = NULL;
  }
  return hr;
}

HRESULT HrReadRaw(IStream *pistm, DWORD *pdw)
{
  HRESULT hr = S_OK;
  DWORD cbSrc = 0;
  DWORD cbRead = 0;

  if (pistm == NULL || pdw == NULL)
    return E_INVALIDARG;

  *pdw = 0;

  hr = pistm->Read((void*)&cbSrc, sizeof(DWORD), &cbRead);
  if (SUCCEEDED(hr) && sizeof(DWORD) != cbRead)
    hr = STG_E_READFAULT;
  if (cbSrc != sizeof(DWORD))
    hr = STG_E_READFAULT;
  if (FAILED(hr))
    goto LExit;

  hr = pistm->Read((void*)pdw, cbSrc, NULL);
  if (FAILED(hr))
    goto LExit;

LExit:
  return hr;
}

HRESULT HrWriteSubStream(IStorage *pistg, const WCHAR *wzStream, BSTR *pbstr)
{
  HRESULT hr = S_OK;
  IStream *pistmT = NULL;

  if (pbstr == NULL)
  {
    hr = E_INVALIDARG;
    goto LExit;
  }

  hr = HrEnsureStm(pistg, wzStream, true /*fReadWrite*/, &pistmT);
  if (FAILED(hr))
    goto LExit;
  if (*pbstr == NULL)
    goto LExit;
  hr = HrWriteRaw(pistmT, *pbstr);
  if (FAILED(hr))
    goto LExit;

LExit:
  if (pistmT != NULL)
    pistmT->Release();
  if (*pbstr != NULL)
  {
    SysFreeString(*pbstr);
    *pbstr = NULL;
  }
  return hr;
}

HRESULT HrReadSubStream(IStorage *pistg, const WCHAR *wzStream, BSTR *pbstr)
{
  HRESULT hr = S_OK;
  IStream *pistmT = NULL;
  WCHAR *wzT = NULL;

  if (pbstr == NULL)
  {
    hr = E_INVALIDARG;
    goto LExit;
  }

  hr = HrEnsureStm(pistg, wzStream, false /*fReadWrite*/, &pistmT);

  if (FAILED(hr))
    goto LExit;	
  hr = HrReadRaw(pistmT, &wzT);
  if (FAILED(hr))
    goto LExit;
  *pbstr = SysAllocString(wzT);
  if (*pbstr == NULL)
  {
    hr = E_OUTOFMEMORY;
    goto LExit;
  }

LExit:
  if (wzT != NULL)
    delete wzT;
  if (pistmT != NULL)
    pistmT->Release();
  return hr;
}

CPdfProtector::CPdfProtector() 
: m_uRefCount(1),
m_langid(1033)
{
  ::InterlockedIncrement(&g_cLocks);
}

CPdfProtector::~CPdfProtector() 
{
  ::InterlockedDecrement(&g_cLocks);
}

HRESULT CPdfProtector::QueryInterface(const IID& riid, void** ppvObj)
{
  if (riid == IID_IUnknown || riid == IID_I_IrmProtector)
  {
    *ppvObj = (I_IrmProtector*)this;
    AddRef();
    return S_OK;
  }
  *ppvObj = 0;
  return E_NOINTERFACE;
}

unsigned long CPdfProtector::AddRef()
{
  return ++m_uRefCount;
}

unsigned long CPdfProtector::Release()
{
  if (--m_uRefCount != 0)
    return m_uRefCount;
  delete this;
  return 0;
}

STDMETHODIMP CPdfProtector::HrInit(BSTR *pbstrProduct, DWORD *pdwVersion, BSTR *pbstrExtensions, BOOL *pfUseRMS)
{
  if ( pbstrExtensions == NULL || pfUseRMS == NULL)
    return E_INVALIDARG;

  *pfUseRMS        = true; 
  *pbstrProduct    = SysAllocString(L"PdfIrmProtector");
  *pdwVersion      = 1;
  *pbstrExtensions = SysAllocString(L"pdf");

  return S_OK;
}

HRESULT CPdfProtector::HrIsProtected(ILockBytes *pilbInput, DWORD *pdwResult)
{
  HRESULT hr = S_OK;
  IStorage *pistgInput = NULL;
  IStream *pistmKey = NULL;
  IStream  *pistmContent = NULL;

  if (pilbInput == NULL || pdwResult == NULL)
  {
    hr = E_INVALIDARG;
    goto LExit;
  }

  *pdwResult = MSOIPI_RESULT_UNKNOWN;
  hr = StgIsStorageILockBytes(pilbInput);
  if (FAILED(hr))
    goto LExit;
  if (hr == S_FALSE)
  {		
    *pdwResult = MSOIPI_RESULT_UNPROTECTED;
    goto LExit;
  }
  hr = StgOpenStorageOnILockBytes(pilbInput, NULL, STGM_READ|STGM_SHARE_EXCLUSIVE, NULL, NULL, &pistgInput);
  if (FAILED(hr))
    goto LExit;
  hr = HrEnsureStm(pistgInput, wzKey, FALSE /*fReadWrite*/, &pistmKey);
  if (FAILED(hr))
  {		
    *pdwResult = MSOIPI_RESULT_UNPROTECTED;
    goto LExit;
  }
  hr = HrEnsureStm(pistgInput, wzContent, FALSE /*fReadWrite*/, &pistmContent);
  if (FAILED(hr))
    
    goto LExit;
  *pdwResult = MSOIPI_RESULT_PROTECTED;

LExit:
  if (pistmContent != NULL)
    pistmContent->Release();
  if (pistmKey != NULL)
    pistmKey->Release();
  if (pistgInput != NULL)
    pistgInput->Release();
  return S_OK;
}

STDMETHODIMP CPdfProtector::HrSetLangId(LANGID langid)
{
  m_langid = langid;
  return S_OK;
}

HRESULT CPdfProtector::HrProtectRMS(ILockBytes *pilbInput, ILockBytes *pilbOutput, I_IrmPolicyInfoRMS *piid, DWORD *pdwStatus)
{
	return E_NOTIMPL;
}

HRESULT CPdfProtector::HrUnprotectRMS(ILockBytes *pilbInput, ILockBytes *pilbOutput, I_IrmPolicyInfoRMS *piid, DWORD *pdwStatus)
{
	return E_NOTIMPL;
}

HRESULT CPdfProtector::HrProtect(ILockBytes *pilbInput, ILockBytes *pilbOutput, I_IrmPolicyInfo *piid, DWORD *pdwStatus)
{
  HRESULT hr = S_OK;
  IStorage *pistgOutput = NULL;
  IStream *pistmT = NULL;
  BSTR bstr = NULL;
  DWORD dw = 0;
  BOOL bRequestingUserIsSystem = false;
  BYTE rgbBuffer[STACK_RW_BUF_SIZE];
  ULARGE_INTEGER ulOffset = { 0 };

  if (pilbInput == NULL || pilbOutput == NULL || piid == NULL || pdwStatus == NULL)
  {
    hr = E_INVALIDARG;
    goto LExit;
  }

  *pdwStatus = MSOIPI_STATUS_UNKNOWN;

  hr = StgIsStorageILockBytes(pilbOutput);
  if (FAILED(hr))
    goto LExit;
  if (hr != S_FALSE)
  {
    hr = StgOpenStorageOnILockBytes(pilbOutput, NULL, STGM_READWRITE|STGM_SHARE_EXCLUSIVE, NULL, NULL, &pistgOutput);
    if (hr == STG_E_FILEALREADYEXISTS)
      goto LCreateStorage;
  }
  else
  {
LCreateStorage:
    hr = StgCreateDocfileOnILockBytes(pilbOutput, STGM_CREATE|STGM_READWRITE|STGM_SHARE_EXCLUSIVE, NULL, &pistgOutput);
  }
  if (FAILED(hr))
    goto LExit;
  BSTR key = SysAllocString(wzKey);
  hr = HrWriteSubStream(pistgOutput, wzKey, &key);
  hr = piid->HrGetListGuid(&bstr);
  if (FAILED(hr) || bstr == NULL)
    goto LExit;
  hr = HrWriteSubStream(pistgOutput, wzListGuid, &bstr);
  if (FAILED(hr))
    goto LExit;
  hr = piid->HrGetRightsMask(&dw);
  if (FAILED(hr))
    goto LExit;
  hr = HrWriteSubStream(pistgOutput, wzRightsMask, &bstr);
  if (FAILED(hr))
    goto LExit;

  hr = piid->HrGetRequestingUser(&bstr, &bRequestingUserIsSystem);
  if (FAILED(hr) || bstr == NULL)
    goto LExit;
  hr = HrWriteSubStream(pistgOutput, wzRequestingUser, &bstr);
  if (FAILED(hr))
    goto LExit;

  hr = piid->HrGetURL(&bstr);
  if (FAILED(hr) || bstr == NULL)
    goto LExit;
  hr = HrWriteSubStream(pistgOutput, wzURL, &bstr);
  if (FAILED(hr))
    goto LExit;

  hr = piid->HrGetPolicyTitle(&bstr);
  if (FAILED(hr) || bstr == NULL)
    goto LExit;
  hr = HrWriteSubStream(pistgOutput, wzPolicyTitle, &bstr);
  if (FAILED(hr))
    goto LExit;

  hr = piid->HrGetPolicyTitle(&bstr);
  if (FAILED(hr) || bstr == NULL)
    goto LExit;
  hr = HrWriteSubStream(pistgOutput, wzPolicyDescription, &bstr);
  if (FAILED(hr))
    goto LExit;

  hr = piid->HrGetOfflineDays(&dw);
  if (FAILED(hr))
    goto LExit;
  hr = HrWriteSubStream(pistgOutput, wzOfflineDays, &bstr);
  if (FAILED(hr))
    goto LExit;

  hr = HrEnsureStm(pistgOutput, wzContent, true /*fReadWrite*/, &pistmT);
  if (FAILED(hr))
    goto LExit;
  ulOffset.QuadPart = 0;

  TCHAR temp_path[ MAX_PATH ];
  const DWORD temp_path_len = GetTempPath( MAX_PATH, temp_path );
  if (!temp_path_len)
	  goto LExit;
  GetTempFileName( temp_path, _T("irm"), FALSE, temp_path );

  HANDLE hTempFile = CreateFile( temp_path, //file name 
    GENERIC_WRITE,  // open for write
    0,              // do not share
    NULL,           // default security
    CREATE_ALWAYS,  // overwrite existing
    FILE_ATTRIBUTE_NORMAL,
    NULL );
  if ( hTempFile == INVALID_HANDLE_VALUE )
    goto LExit;

  while (true)
  {
    ULONG cbRead = cElements(rgbBuffer);
    ULONG cbWritten = 0;
    hr = pilbInput->ReadAt(ulOffset, rgbBuffer, cbRead, &cbRead);
    if (FAILED(hr))
      goto LExit;
    if (cbRead == 0)
      goto LFinished;

    WriteFile( hTempFile, rgbBuffer, cbRead, &cbWritten, NULL );

    if (cbRead != cbWritten)
    {
      hr = STG_E_CANTSAVE;
      goto LExit;
    }
    ulOffset.QuadPart += cbRead;
  }

LFinished:
  TCHAR temp_path2[ MAX_PATH ];
  const DWORD temp_path2_len = GetTempPath( MAX_PATH, temp_path2 );
  if (!temp_path2_len)
	  goto LExit;
  GetTempFileName( temp_path2, _T("irm"), FALSE, temp_path2 );

  CloseHandle( hTempFile );
  {
    const char* password = NULL;//"text";
    const char* user_password = NULL;
    const bool decrypt = false;
    const bool encrypt = true;

    USES_CONVERSION;
    QPDF pdf;
    pdf.processFile(T2A(temp_path), password);
	  QPDFWriter w(pdf, T2A(temp_path2));
    if (decrypt)
	  {
		  w.setPreserveEncryption(false);
	  }
    if (encrypt)
    {
      w.setR3EncryptionParameters(
        user_password, W2A(wzKey),
        false, //r3_accessibility
        false, //r3_extract 
        qpdf_r3p_none, //r3_print 
        qpdf_r3m_none  //r3_modify
        );
    }
    w.write();

    HANDLE hOutFile = CreateFile( temp_path2, GENERIC_READ, 0, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL );
    if ( hOutFile == INVALID_HANDLE_VALUE )
    {
      goto LExit;
    }

    DWORD cbReaded = 0, cbWritten;
    while ( ReadFile( hOutFile, rgbBuffer, sizeof( rgbBuffer ), &cbReaded, NULL ) || cbReaded )
    {
      hr = pistmT->Write(rgbBuffer, cbReaded, &cbWritten);
      if (FAILED(hr))
        goto LExit;
      if ( cbReaded != cbWritten )
      {
        hr = STG_E_CANTSAVE;
        goto LExit;
      }
    }
  }

  hr = pistmT->Commit(STGC_DEFAULT);

LExit:
  DeleteFile( temp_path );
  DeleteFile( temp_path2 );

  *pdwStatus = SUCCEEDED(hr) ? MSOIPI_STATUS_PROTECT_SUCCESS : MSOIPI_STATUS_CANT_PROTECT;
  if (bstr != NULL)
    SysFreeString(bstr);
  if (pistmT != NULL)
    pistmT->Release();
  if (pistgOutput != NULL)
    pistgOutput->Release();

  return hr;
}

HRESULT CPdfProtector::HrUnprotect(ILockBytes *pilbInput, ILockBytes *pilbOutput, I_IrmPolicyInfo *piid, DWORD *pdwStatus)
{
  HRESULT hr = S_OK;
  IStorage *pistgInput = NULL;
  IStream *pistmT = NULL;
  BYTE rgbBuffer[STACK_RW_BUF_SIZE];
  ULARGE_INTEGER ulOffset = { 0 };

  if (pilbInput == NULL || pilbOutput == NULL || piid == NULL || pdwStatus == NULL)
  {
    return E_INVALIDARG;
  }

  *pdwStatus = MSOIPI_STATUS_UNKNOWN;

  hr = StgIsStorageILockBytes(pilbInput);
  if (FAILED(hr) || hr == S_FALSE)
  {
    *pdwStatus = MSOIPI_STATUS_ALREADY_UNPROTECTED;
    goto LExit;
  }

  hr = StgOpenStorageOnILockBytes(pilbInput, NULL, STGM_READ|STGM_SHARE_EXCLUSIVE, NULL, NULL, &pistgInput);
  if (FAILED(hr))
    goto LExit;

  hr = HrEnsureStm(pistgInput, wzContent, false /*fReadWrite*/, &pistmT);
  if (FAILED(hr))
    goto LExit;
  ulOffset.QuadPart = 0;

  while (true)
  {
    ULONG cbRead = cElements(rgbBuffer);
    ULONG cbWritten = 0;
    hr = pistmT->Read(rgbBuffer, cbRead, &cbRead);
    if (FAILED(hr))
      goto LExit;
    if (cbRead == 0)
      goto LFinished;
    for (ULONG i = 0; i < cbRead; i++)
      //Flip the bits back
      rgbBuffer[i] ^= 1;
    hr = pilbOutput->WriteAt(ulOffset, rgbBuffer, cbRead, &cbWritten);
    if (FAILED(hr))
      goto LExit;
    if (cbRead != cbWritten)
    {
      hr = STG_E_CANTSAVE;
      goto LExit;
    }
    ulOffset.QuadPart += cbRead;
  }
LFinished:
  pistmT->Release();
  pistmT = NULL;

LExit:
  if (pistmT != NULL)
    pistmT->Release();
  if (pistgInput != NULL)
    pistgInput->Release();
  return hr;
}